"""
基于YOLO-OBB的旋转目标检测模块
"""

import cv2
import numpy as np
import torch
import logging
from typing import List, Dict, Any

class YOLO_OBB_Detector:
    """YOLO-OBB旋转目标检测器"""
    
    def __init__(self, config):
        self.config = config
        self.logger = logging.getLogger(__name__)
        self.device = self.setup_device()
        self.model = self.load_model()
        
    def setup_device(self):
        """设置计算设备"""
        if self.config.use_gpu and torch.cuda.is_available():
            device = torch.device('cuda')
            self.logger.info("使用GPU进行检测")
        else:
            device = torch.device('cpu')
            self.logger.info("使用CPU进行检测")
        return device
    
    def load_model(self):
        """加载YOLO-OBB模型"""
        try:
            # 这里使用Ultralytics YOLO
            from ultralytics import YOLO
            model = YOLO(self.config.model_path)
            self.logger.info(f"成功加载模型: {self.config.model_path}")
            return model
        except Exception as e:
            self.logger.error(f"模型加载失败: {str(e)}")
            raise
    
    def preprocess(self, image):
        """图像预处理"""
        # 调整图像尺寸
        input_size = self.config.input_size
        h, w = image.shape[:2]
        
        # 保持宽高比的resize
        scale = min(input_size / w, input_size / h)
        new_w, new_h = int(w * scale), int(h * scale)
        
        resized = cv2.resize(image, (new_w, new_h))
        
        # 填充到正方形
        padded = np.full((input_size, input_size, 3), 114, dtype=np.uint8)
        padded[:new_h, :new_w] = resized
        
        # 归一化
        padded = padded.astype(np.float32) / 255.0
        
        # 转换为CHW格式
        padded = padded.transpose(2, 0, 1)
        
        return padded, scale, (w, h)
    
    def postprocess(self, predictions, scale, orig_shape):
        """后处理检测结果"""
        detections = []
        
        if predictions is None or len(predictions) == 0:
            return detections
            
        for pred in predictions[0].obb:  # OBB预测结果
            if pred.confidence < self.config.confidence_threshold:
                continue
                
            # 解析旋转框参数
            xc, yc, w, h, angle = pred.xywhr[0].tolist()
            
            # 缩放回原图尺寸
            xc = xc / scale
            yc = yc / scale
            w = w / scale
            h = h / scale
            
            # 计算旋转框的四个角点
            corners = self.calculate_rotated_corners(xc, yc, w, h, angle)
            
            detection = {
                'bbox': [xc, yc, w, h, angle],
                'corners': corners,
                'confidence': pred.confidence,
                'class_id': pred.cls,
                'class_name': pred.names[int(pred.cls)] if hasattr(pred, 'names') else 'vehicle'
            }
            
            detections.append(detection)
            
        # 非极大值抑制
        detections = self.non_max_suppression(detections)
        
        return detections
    
    def calculate_rotated_corners(self, xc, yc, w, h, angle):
        """计算旋转矩形的四个角点"""
        cos_a = np.cos(angle)
        sin_a = np.sin(angle)
        
        # 半宽半高
        half_w = w / 2
        half_h = h / 2
        
        # 四个角点相对于中心的坐标
        corners_rel = np.array([
            [-half_w, -half_h],
            [half_w, -half_h],
            [half_w, half_h],
            [-half_w, half_h]
        ])
        
        # 旋转
        rotation_matrix = np.array([
            [cos_a, -sin_a],
            [sin_a, cos_a]
        ])
        
        corners_rotated = corners_rel @ rotation_matrix.T
        
        # 平移回图像坐标
        corners_absolute = corners_rotated + np.array([xc, yc])
        
        return corners_absolute.tolist()
    
    def non_max_suppression(self, detections):
        """非极大值抑制"""
        if len(detections) <= 1:
            return detections
            
        # 按置信度排序
        detections.sort(key=lambda x: x['confidence'], reverse=True)
        
        keep = []
        while detections:
            # 取置信度最高的检测框
            best = detections.pop(0)
            keep.append(best)
            
            # 计算与剩余框的IoU
            to_remove = []
            for i, det in enumerate(detections):
                iou = self.calculate_rotated_iou(best['corners'], det['corners'])
                if iou > self.config.nms_threshold:
                    to_remove.append(i)
            
            # 移除重叠框
            for i in sorted(to_remove, reverse=True):
                detections.pop(i)
                
        return keep
    
    def calculate_rotated_iou(self, corners1, corners2):
        """计算旋转框IoU"""
        try:
            # 使用OpenCV的旋转矩形IoU计算
            rect1 = cv2.minAreaRect(np.array(corners1, dtype=np.float32))
            rect2 = cv2.minAreaRect(np.array(corners2, dtype=np.float32))
            
            # 计算IoU
            inter_area, _ = cv2.rotatedRectangleIntersection(rect1, rect2)
            if inter_area == 0:
                return 0.0
                
            area1 = cv2.contourArea(cv2.boxPoints(rect1))
            area2 = cv2.contourArea(cv2.boxPoints(rect2))
            
            iou = inter_area / (area1 + area2 - inter_area)
            return iou
            
        except Exception as e:
            self.logger.warning(f"IoU计算错误: {str(e)}")
            return 0.0
    
    def detect(self, image):
        """执行检测"""
        try:
            # 预处理
            processed, scale, orig_shape = self.preprocess(image)
            
            # 推理
            with torch.no_grad():
                results = self.model(processed, imgsz=self.config.input_size)
            
            # 后处理
            detections = self.postprocess(results, scale, orig_shape)
            
            return detections[:self.config.max_detections]
            
        except Exception as e:
            self.logger.error(f"检测错误: {str(e)}")
            return []